package top.longmarch.flow.ai;

import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.json.JSONUtil;
import okhttp3.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;


public class AIUtil {
    private static final Logger log = LoggerFactory.getLogger(AIUtil.class);
    private static final ExecutorService executorService = ThreadUtil.newExecutor(10);
    public static final String AI_URL = "http://180.184.214.46:31398/openapi/llm/chat";
    public static final String API_KEY = "sk-3bddd74413064722b831a9429e3588fd";
    public static final String MODEL = "qwen-plus";

    public static final String PROMPT1 = "你是一个专业的代码解释专家，请根据输入的代码，提炼中文名称，请求数据对象（this.getRequestData），上下文对象（this.getContextBean），依赖对象。再给出代码的解释，中文回复。";

    public static Request getOkHttpRequest(String content) {
        return new Request.Builder()
                .url(AIUtil.AI_URL)
                .header("Authorization", "Bearer " + AIUtil.API_KEY)
                .post(okhttp3.RequestBody.create(getParams(content, true), MediaType.parse("application/json")))
                .build();
    }

    public static String getParams(String content, Boolean stream) {
        Map<String, Object> params = new HashMap<>();
        params.put("model", MODEL);
        params.put("temperature", 0.7);
        params.put("stream", stream);

        List<Map<String, Object>> messages = new ArrayList<>();
        Map<String, Object> sysMsg = new HashMap<>();
        sysMsg.put("role", "system");
        sysMsg.put("content", PROMPT1);
        messages.add(sysMsg);

        Map<String, Object> userMsg = new HashMap<>();
        userMsg.put("role", "user");
        userMsg.put("content", content);
        messages.add(userMsg);
        params.put("messages", messages);
        return JSONUtil.toJsonStr(params);
    }

    public static ResponseBodyEmitter streamChat2(String content) {
        SseEmitter emitter = AsyncResponseUtil.createUtf8SseEmitter();
        AsyncResponseUtil.WrappedResponseBodyEmitter wrapped = AsyncResponseUtil.wrapEmitter(emitter);
        CompletableFuture.runAsync(() -> {
            try {
                processRequest(content, emitter, wrapped);
            } catch (Exception e) {
                log.error("Async task failed", e);
                AsyncResponseUtil.complete(emitter);
            }
        }, executorService).exceptionally(e -> {
            log.error("Exception in CompletableFuture", e);
            return null;
        }).whenComplete((t, e) -> {
            try {
                AsyncResponseUtil.complete(emitter);
            } catch (Exception exception) {
                log.error("Failed to complete emitter", exception);
            }
        });
        return emitter;
    }

    private static void processRequest(String content, SseEmitter emitter, AsyncResponseUtil.WrappedResponseBodyEmitter wrapped) throws Exception {
        Request request = getOkHttpRequest(content);
        OkHttpClient client = new OkHttpClient();
        try (Response response = client.newCall(request).execute()) {
            if (!response.isSuccessful()) {
                log.error("HTTP failed: " + response.code());
                throw new IOException("HTTP request failed with code: " + response.code());
            }
            ResponseBody body = response.body();
            try (InputStreamReader ise = new InputStreamReader(body.byteStream(), StandardCharsets.UTF_8);
                 BufferedReader reader = new BufferedReader(ise)) {
                reader.lines().forEach(line -> {
                    int colonIndex = line.indexOf(':');
                    if (colonIndex > 0) {
                        String field = line.substring(0, colonIndex);
                        String value = line.substring(colonIndex + 1);
                        if ("data".equals(field) && !wrapped.isCompleted()) {
                            AsyncResponseUtil.send(emitter, value);
                        }
                    }
                });
            }
        }
    }

}
